import os
import yaml
import shutil
from datetime import datetime, timezone
from pathlib import Path

# Metadata Libraries
from mutagen.wave import WAVE
from mutagen.mp4 import MP4
from mutagen.id3 import ID3, TIT2, TPE1, TALB, TCOP, COMM
from pikepdf import Pdf, Dictionary, Name
from PIL import Image, PngImagePlugin

# --- CONFIGURATION ---
OUTPUT_FOLDER_NAME = "Branded_Public_Domain_Assets"

def get_archival_timestamp():
    """Generates the archival UTC timestamp."""
    return datetime.now(timezone.utc).strftime('%Y-%m-%dT%H:%M:%SZ')

def process_wav(input_path, output_path, meta, ts):
    """Handles WAV metadata (RIFF Tags)."""
    shutil.copy2(input_path, output_path)
    audio = WAVE(output_path)
    # WAV uses RIFF tags: IPRD is Album/Product, IART is Artist, ICOP is Copyright
    audio.tags['IPRD'] = meta.get('url', '') # Album field
    audio.tags['IART'] = meta.get('author', '')
    audio.tags['ICOP'] = meta.get('license', '')
    audio.tags['ICMT'] = meta.get('comment', '').replace("[ARCHIVAL TIMESTAMP]", ts)
    audio.tags['ISBJ'] = meta.get('brand', '')
    audio.save()

def process_pdf(input_path, output_path, meta, ts):
    """Handles PDF metadata."""
    with Pdf.open(input_path) as pdf:
        with pdf.open_metadata() as pdf_meta:
            pdf_meta['dc:title'] = input_path.stem
            pdf_meta['dc:creator'] = [meta.get('author', '')]
            pdf_meta['dc:description'] = meta.get('comment', '').replace("[ARCHIVAL TIMESTAMP]", ts)
            pdf_meta['dc:rights'] = meta.get('license', '')
            pdf_meta['xmp:CreatorTool'] = meta.get('brand', '')
            # PDFs don't have an 'Album' field, so we use a custom property or Subject
            pdf_meta['pdf:Keywords'] = meta.get('url', '')
        pdf.save(output_path)

def process_image(input_path, output_path, meta, ts):
    """Handles Image metadata (JPG/PNG)."""
    img = Image.open(input_path)
    comment = meta.get('comment', '').replace("[ARCHIVAL TIMESTAMP]", ts)
    
    # Copy the image to maintain quality, then save with metadata
    if input_path.suffix.lower() in ['.jpg', '.jpeg']:
        # For JPEG, we use the 'comment' and basic info
        exif = img.getexif()
        exif[0x0131] = meta.get('author', '') # Software/Artist
        exif[0x8298] = meta.get('license', '') # Copyright
        exif[0x9286] = comment # User Comment
        img.save(output_path, exif=exif, quality='keep')
    else:
        # For PNG
        png_meta = PngImagePlugin.PngInfo()
        png_meta.add_text("Author", meta.get('author', ''))
        png_meta.add_text("Copyright", meta.get('license', ''))
        png_meta.add_text("Comment", comment)
        png_meta.add_text("URL", meta.get('url', ''))
        img.save(output_path, pnginfo=png_meta)

def process_video(input_path, output_path, meta, ts):
    """Handles MP4 Video metadata."""
    shutil.copy2(input_path, output_path)
    try:
        video = MP4(output_path)
        video["\xa9nam"] = input_path.stem # Title
        video["\xa9ART"] = meta.get('author', '') # Artist
        video["\xa9alb"] = meta.get('url', '') # Album field used for URL
        video["\xa9cmt"] = meta.get('comment', '').replace("[ARCHIVAL TIMESTAMP]", ts)
        video["cprt"] = meta.get('license', '')
        video.save()
    except Exception as e:
        print(f"      Skipping advanced tags for {input_path.name}: {e}")

def main():
    current_dir = Path('.')
    ts = get_archival_timestamp()

    for folder in [d for d in current_dir.iterdir() if d.is_dir()]:
        if folder.name == OUTPUT_FOLDER_NAME: continue
        
        yaml_file = folder / "metadata.yaml"
        if not yaml_file.exists(): continue

        print(f"\n📂 Processing: {folder.name}")
        
        with open(yaml_file, 'r', encoding='utf-8') as f:
            meta = yaml.safe_load(f)

        out_path = folder / OUTPUT_FOLDER_NAME
        out_path.mkdir(exist_ok=True)

        for file in folder.iterdir():
            ext = file.suffix.lower()
            dest = out_path / file.name
            
            try:
                if ext == '.wav':
                    print(f"  🏷️  Audio: {file.name}")
                    process_wav(file, dest, meta, ts)
                elif ext == '.pdf':
                    print(f"  🏷️  PDF: {file.name}")
                    process_pdf(file, dest, meta, ts)
                elif ext in ['.jpg', '.jpeg', '.png']:
                    print(f"  🏷️  Image: {file.name}")
                    process_image(file, dest, meta, ts)
                elif ext in ['.mp4', '.m4v']:
                    print(f"  🏷️  Video: {file.name}")
                    process_video(file, dest, meta, ts)
            except Exception as e:
                print(f"  ❌ Failed {file.name}: {e}")

if __name__ == "__main__":
    main()
    print("\n✅ Branded Assets Complete.")